package com.yuan.aiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

@Configuration
public class PgVectorVectorStoreConfig {
    @Resource
    private AppDocumentLoader documentLoader;
    @Resource
    private MyDocumentEnricher documentEnricher;

    @Primary
    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, @Qualifier("ollamaEmbeddingModel")EmbeddingModel embeddingModel) {
        PgVectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, embeddingModel)
                .dimensions(1536)
                .distanceType(COSINE_DISTANCE)
                .indexType(HNSW)
                .schemaName("ai_vector")
                .vectorTableName("vector_store")
                .maxDocumentBatchSize(10000)
                .build();
        List markdowns = documentLoader.loadMarkdowns();
        List enrichedDocuments = documentEnricher.enrichDocumentsByKeyword(markdowns);
        vectorStore.add(enrichedDocuments);
        return vectorStore;
    }
}
